import os
import csv
import numpy as np
import joblib
from sample_predictor import predict_sample
from asm_feature_extractor import load_labels, process_asm_file

# 配置文件路径
MODEL_PATH = "models/codebert-base"
ASM_DIR = "samples/subtrain/selectedpro/test"  # ASM测试数据目录
LABEL_FILE = "samples/selectedLabels.csv"
CUSTOM_FILE = "customs/codebert_selectedpro_classifier.joblib"  # 训练完成的模型文件
PREDICTION_RESULTS = "results/prediction_results.csv"  # 预测结果文件

def predict_all_samples():
    """使用训练好的模型预测所有样本"""
    # 加载标签
    true_labels = load_labels(LABEL_FILE)
    
    # 加载模型
    if not os.path.exists(CUSTOM_FILE):
        print(f"错误: 模型文件 {CUSTOM_FILE} 不存在!")
        return
    
    print(f"加载模型: {CUSTOM_FILE}")
    pipeline = joblib.load(CUSTOM_FILE)
    
    # 获取所有样本ID
    file_ids = [f.replace('.asm', '') for f in os.listdir(ASM_DIR) if f.endswith('.asm')]
    
    # 准备结果存储
    predictions = []
    probabilities = []
    correct_count = 0
    total_count = 0
    
    # 创建结果CSV文件
    os.makedirs(os.path.dirname(PREDICTION_RESULTS), exist_ok=True)
    with open(PREDICTION_RESULTS, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['File_ID', 'True_Class', 'Predicted_Class', 'Probability', 'Correct', 
                      'Prob_Class1', 'Prob_Class2', 'Prob_Class3', 'Prob_Class4', 'Prob_Class5',
                      'Prob_Class6', 'Prob_Class7', 'Prob_Class8', 'Prob_Class9']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        
        print(f"开始预测 {len(file_ids)} 个样本...")
        
        for i, file_id in enumerate(file_ids):
            if file_id not in true_labels:
                continue
                
            asm_path = os.path.join(ASM_DIR, file_id + '.asm')
            print(f"[{i+1}/{len(file_ids)}] 预测: {file_id}")
            
            try:
                # 使用predict_sample接口进行预测
                prediction, probability = predict_sample(MODEL_PATH, asm_path, pipeline)
                print(f"最可能的类别: {prediction}")
                
                if prediction is None:
                    print(f"  ⚠️ 无法预测: {file_id}")
                    continue
                
                # 获取真实标签
                true_label = true_labels[file_id]
                
                # 检查是否正确
                is_correct = int(prediction == true_label)
                if is_correct:
                    correct_count += 1
                total_count += 1
                
                # 获取概率分布（需要从predict_sample中获取）
                # 注意：这里需要修改predict_sample函数以返回概率分布
                # 假设predict_sample返回一个元组 (prediction, probability)
                # prediction, probability = predict_sample(MODEL_PATH, asm_path, pipeline)
                
                # 由于predict_sample接口没有返回概率，我们需要单独获取
                # 这里我们直接使用pipeline获取概率
                feature_vector = process_asm_file(MODEL_PATH, asm_path)
                if feature_vector is None or len(feature_vector) == 0:
                    print(f"  ⚠️ 无法提取特征: {file_id}")
                    continue
                
                # 修复：使用安全的形状转换方法
                if isinstance(feature_vector, (list, tuple)):
                    feature_vector = np.array(feature_vector)
                
                # 转换为正确的形状 (1, n_features)
                if feature_vector.ndim == 1:
                    feature_vector = feature_vector.reshape(1, -1)
                
                # 获取概率分布
                if hasattr(pipeline, 'predict_proba'):
                    probability = pipeline.predict_proba(feature_vector)[0]
                else:
                    # 对于不支持概率预测的模型，使用决策函数
                    decision_values = pipeline.decision_function(feature_vector)[0]
                    probability = np.exp(decision_values) / np.sum(np.exp(decision_values))
                
                # 存储结果
                predictions.append(prediction)
                probabilities.append(probability)
                
                # 写入CSV
                writer.writerow({
                    'File_ID': file_id,
                    'True_Class': true_label,
                    'Predicted_Class': prediction,
                    'Probability': probability[prediction-1] if prediction <= len(probability) else 0,
                    'Correct': is_correct,
                    'Prob_Class1': probability[0] if len(probability) > 0 else 0,
                    'Prob_Class2': probability[1] if len(probability) > 1 else 0,
                    'Prob_Class3': probability[2] if len(probability) > 2 else 0,
                    'Prob_Class4': probability[3] if len(probability) > 3 else 0,
                    'Prob_Class5': probability[4] if len(probability) > 4 else 0,
                    'Prob_Class6': probability[5] if len(probability) > 5 else 0,
                    'Prob_Class7': probability[6] if len(probability) > 6 else 0,
                    'Prob_Class8': probability[7] if len(probability) > 7 else 0,
                    'Prob_Class9': probability[8] if len(probability) > 8 else 0
                })
                
            except Exception as e:
                print(f"  ❌ 预测失败: {file_id} - {str(e)}")
    
    # 计算准确率
    if total_count > 0:
        accuracy = correct_count / total_count
        print("\n" + "="*50)
        print(f"预测完成! 总样本数: {total_count}, 正确预测: {correct_count}")
        print(f"整体准确率: {accuracy:.4f} ({accuracy*100:.2f}%)")
        print(f"结果已保存至: {PREDICTION_RESULTS}")
        print("="*50)
    else:
        print("没有成功预测任何样本")

if __name__ == "__main__":
    predict_all_samples()